#!/bin/bash

# K8s集群主机名配置脚本
# 自动配置所有节点的主机名和hosts文件

set -e

# 配置文件路径
CONFIG_FILE="./k8s-config.conf"
LOG_FILE="/var/log/k8s-hostname.log"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 日志函数
log() {
    local level=$1
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case $level in
        "INFO")
            echo -e "${BLUE}[INFO]${NC} $message" | tee -a $LOG_FILE
            ;;
        "SUCCESS")
            echo -e "${GREEN}[SUCCESS]${NC} $message" | tee -a $LOG_FILE
            ;;
        "WARNING")
            echo -e "${YELLOW}[WARNING]${NC} $message" | tee -a $LOG_FILE
            ;;
        "ERROR")
            echo -e "${RED}[ERROR]${NC} $message" | tee -a $LOG_FILE
            ;;
    esac
}

# 错误处理函数
error_exit() {
    log "ERROR" "$1"
    log "ERROR" "主机名配置失败，请检查日志: $LOG_FILE"
    exit 1
}

# 检查配置文件
check_config() {
    if [ ! -f "$CONFIG_FILE" ]; then
        error_exit "配置文件不存在: $CONFIG_FILE"
    fi
    
    # 加载配置文件
    source "$CONFIG_FILE"
    
    log "SUCCESS" "配置文件加载成功"
    log "INFO" "Master节点: $MASTER_HOSTNAME ($MASTER_IP)"
    log "INFO" "Worker1节点: $WORKER1_HOSTNAME ($WORKER1_IP)"
    log "INFO" "Worker2节点: $WORKER2_HOSTNAME ($WORKER2_IP)"
}

# 检查root权限
check_root() {
    if [ "$EUID" -ne 0 ]; then
        error_exit "请使用root用户运行此脚本"
    fi
    log "SUCCESS" "权限检查通过"
}

# 配置本地主机名
configure_local_hostname() {
    local node_type=$1
    local hostname=""
    local ip=""
    
    case $node_type in
        "master")
            hostname="$MASTER_HOSTNAME"
            ip="$MASTER_IP"
            ;;
        "worker1")
            hostname="$WORKER1_HOSTNAME"
            ip="$WORKER1_IP"
            ;;
        "worker2")
            hostname="$WORKER2_HOSTNAME"
            ip="$WORKER2_IP"
            ;;
        *)
            error_exit "无效的节点类型: $node_type"
            ;;
    esac
    
    log "INFO" "配置本地主机名: $hostname"
    
    # 设置主机名
    hostnamectl set-hostname "$hostname"
    
    # 更新/etc/hosts
    local current_ip=$(hostname -I | awk '{print $1}')
    
    # 备份原始hosts文件
    cp /etc/hosts /etc/hosts.backup
    
    # 创建新的hosts文件
    cat > /etc/hosts << EOF
# K8s集群主机配置
127.0.0.1 localhost
127.0.1.1 $hostname

# K8s集群节点
$MASTER_IP $MASTER_HOSTNAME
$WORKER1_IP $WORKER1_HOSTNAME
$WORKER2_IP $WORKER2_HOSTNAME

# 以下为系统自动生成的配置
EOF
    
    log "SUCCESS" "本地主机名配置完成: $hostname"
}

# 远程配置主机名
configure_remote_hostname() {
    local hostname=$1
    local ip=$2
    local user=$3
    
    log "INFO" "远程配置主机名: $hostname ($ip)"
    
    # 检查SSH连接
    if ! ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$user@$ip" "echo 'SSH连接成功'" &>/dev/null; then
        log "WARNING" "无法连接到 $ip，跳过配置"
        return 1
    fi
    
    # 远程执行主机名配置
    ssh -o StrictHostKeyChecking=no "$user@$ip" << EOF
        # 设置主机名
        hostnamectl set-hostname $hostname
        
        # 备份原始hosts文件
        cp /etc/hosts /etc/hosts.backup
        
        # 创建新的hosts文件
        cat > /etc/hosts << 'HOSTS_EOF'
# K8s集群主机配置
127.0.0.1 localhost
127.0.1.1 $hostname

# K8s集群节点
$MASTER_IP $MASTER_HOSTNAME
$WORKER1_IP $WORKER1_HOSTNAME
$WORKER2_IP $WORKER2_HOSTNAME

# 以下为系统自动生成的配置
HOSTS_EOF
        
        echo "主机名配置完成: $hostname"
EOF
    
    log "SUCCESS" "远程主机名配置完成: $hostname"
}

# 验证主机名配置
verify_hostnames() {
    log "INFO" "验证主机名配置..."
    
    # 检查本地主机名
    local current_hostname=$(hostname)
    log "INFO" "当前主机名: $current_hostname"
    
    # 检查hosts文件
    log "INFO" "检查hosts文件配置:"
    grep -E "($MASTER_IP|$WORKER1_IP|$WORKER2_IP)" /etc/hosts | while read line; do
        log "INFO" "  $line"
    done
    
    # 测试主机名解析
    log "INFO" "测试主机名解析:"
    for host in $MASTER_HOSTNAME $WORKER1_HOSTNAME $WORKER2_HOSTNAME; do
        if nslookup "$host" &>/dev/null; then
            log "SUCCESS" "✓ $host 解析正常"
        else
            log "WARNING" "✗ $host 解析失败"
        fi
    done
    
    log "SUCCESS" "主机名配置验证完成"
}

# 生成主机名配置报告
generate_report() {
    local report_file="/tmp/k8s-hostname-report-$(date +%Y%m%d-%H%M%S).txt"
    
    log "INFO" "生成主机名配置报告: $report_file"
    
    {
        echo "K8s集群主机名配置报告"
        echo "生成时间: $(date)"
        echo "当前主机: $(hostname)"
        echo "=========================================="
        echo ""
        
        echo "集群节点配置:"
        echo "Master节点: $MASTER_HOSTNAME ($MASTER_IP)"
        echo "Worker1节点: $WORKER1_HOSTNAME ($WORKER1_IP)"
        echo "Worker2节点: $WORKER2_HOSTNAME ($WORKER2_IP)"
        echo ""
        
        echo "当前主机信息:"
        echo "主机名: $(hostname)"
        echo "IP地址: $(hostname -I | awk '{print $1}')"
        echo ""
        
        echo "hosts文件内容:"
        cat /etc/hosts
        echo ""
        
        echo "主机名解析测试:"
        for host in $MASTER_HOSTNAME $WORKER1_HOSTNAME $WORKER2_HOSTNAME; do
            echo -n "$host: "
            if nslookup "$host" &>/dev/null; then
                echo "✓ 解析正常"
            else
                echo "✗ 解析失败"
            fi
        done
        echo ""
        
    } > $report_file
    
    log "SUCCESS" "主机名配置报告已生成: $report_file"
}

# 显示使用说明
show_usage() {
    echo "K8s集群主机名配置脚本"
    echo ""
    echo "用法:"
    echo "  $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --node-type=TYPE     节点类型 (master|worker1|worker2)"
    echo "  --all-nodes          配置所有节点"
    echo "  --verify-only        仅验证配置"
    echo "  --help               显示帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 --node-type=master"
    echo "  $0 --node-type=worker1"
    echo "  $0 --all-nodes"
    echo "  $0 --verify-only"
}

# 主函数
main() {
    log "INFO" "开始K8s集群主机名配置..."
    log "INFO" "配置日志: $LOG_FILE"
    
    # 解析命令行参数
    local node_type=""
    local all_nodes=false
    local verify_only=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            --node-type=*)
                node_type="${1#*=}"
                shift
                ;;
            --all-nodes)
                all_nodes=true
                shift
                ;;
            --verify-only)
                verify_only=true
                shift
                ;;
            --help)
                show_usage
                exit 0
                ;;
            *)
                echo "未知参数: $1"
                show_usage
                exit 1
                ;;
        esac
    done
    
    # 执行配置步骤
    check_root
    check_config
    
    if [ "$verify_only" = true ]; then
        verify_hostnames
        generate_report
        exit 0
    fi
    
    if [ "$all_nodes" = true ]; then
        # 配置所有节点
        log "INFO" "配置所有节点主机名..."
        
        # 配置本地节点（假设当前是master节点）
        configure_local_hostname "master"
        
        # 配置远程节点
        configure_remote_hostname "$WORKER1_HOSTNAME" "$WORKER1_IP" "$WORKER1_USER" || true
        configure_remote_hostname "$WORKER2_HOSTNAME" "$WORKER2_IP" "$WORKER2_USER" || true
        
    elif [ -n "$node_type" ]; then
        # 配置指定节点
        configure_local_hostname "$node_type"
    else
        echo "错误: 请指定节点类型或使用 --all-nodes"
        show_usage
        exit 1
    fi
    
    # 验证配置
    verify_hostnames
    
    # 生成报告
    generate_report
    
    log "SUCCESS" "K8s集群主机名配置完成！"
}

# 运行主函数
main "$@"
